import numpy as np
from spline import Spline


class ReferPath:
    def __init__(self, spline: Spline) -> None:
        self.spline = spline

    def sampler(self, distance: float) -> list[float]:
        """
        distance: float, distance in path
        return: list[float], [theta, v, x, y, s]
        """
        for i in range(len(self.spline.s)):
            if self.spline.s[i] > distance:
                return [
                    self.spline.theta[i],
                    self.spline.v[i],
                    self.spline.x[i],
                    self.spline.y[i],
                    self.spline.s[i],
                ]
        return [
            self.spline.theta[-1],
            self.spline.v[-1],
            self.spline.x[-1],
            self.spline.y[-1],
            self.spline.s[-1],
        ]

    def getNearestPoint(self, pos: tuple[float, float]) -> tuple[int, list[float]]:
        """
        pos: tuple[float,float], [x,y]
        return: tuple[int,list[float]], [index, [theta,v,x,y,s]]
        """
        dx = np.array(self.spline.x) - pos[0]
        dy = np.array(self.spline.y) - pos[1]
        d = dx**2 + dy**2
        index = np.argmin(d)
        return index.item(), [
            self.spline.theta[index],
            self.spline.v[index],
            self.spline.x[index],
            self.spline.y[index],
            self.spline.s[index],
        ]

    def getReferancePath(
        self, pos: tuple[float, float], T: int, timeStep: float
    ) -> tuple[np.ndarray, np.ndarray]:
        index, start = self.getNearestPoint(pos)
        curS = start[4]  # s
        curV = start[1]  # v
        referStates = np.zeros((4, T + 1))
        referControl = np.zeros((2, T))
        for i in range(T + 1):
            states = self.sampler(curS + timeStep * curV)
            curS = states[4]
            curV = states[1]
            referStates[0, i] = states[0]
            referStates[1, i] = states[1]
            referStates[2, i] = states[2]
            referStates[3, i] = states[3]
        for i in range(T):
            if i < T - 1:
                referControl[0, i] = (
                    referStates[0, i + 1] - referStates[0, i]
                ) / timeStep
                referControl[1, i] = (
                    referStates[1, i + 1] - referStates[1, i]
                ) / timeStep
            else:
                referControl[:, i] = referControl[:, i - 1]
        return referStates, referControl

    def drawReferancePath(self, referState: np.ndarray, ax, color: str,index:bool=False,points:bool=False):
        import matplotlib.pyplot as plt

        x = referState[2, :]
        y = referState[3, :]
        ax.plot(x, y, color=color)
        if points:
            ax.scatter(x, y, color=color)
        if index:
            for i in range(len(x)):
                plt.text(x[i], y[i], str(i))


if __name__ == "__main__":
    from spline import ConstVelocitySpline
    from refer_path import ReferPath
    import matplotlib.pyplot as plt

    line = ConstVelocitySpline([(0, 0), (1, 1), (2, 0), (3, -1), (4, 0)], 1, 0.02)
    path = ReferPath(line)
    start = (0.9, 0.5)
    T = 10
    referState, referControl = path.getReferancePath(start, T, 0.1)

    ax = plt.subplot(1, 2, 1)
    line.drawPath(ax, "blue")
    path.drawReferancePath(referState, ax, "green")
    ax.scatter(start[0], start[1], color="red")

    bx = plt.subplot(1, 2, 2)
    bx.plot(referState[0, :], label="theta")
    bx.plot(referState[1, :], label="v")
    bx.plot(referState[2, :], label="x")
    bx.plot(referState[3, :], label="y")
    bx.plot(referControl[0, :], label="w")
    bx.plot(referControl[1, :], label="a")
    bx.legend()

    plt.show()
